import requests
import pandas as pd
from tqdm import tqdm
import pickle

df = pd.read_csv("../Data/ml-1m/movies.dat", sep='::', header=None, encoding="ISO-8859-1")
df.columns = ['mid', 'movie_name', 'classes']
sid2movie_map = {}
for index, row in df.iterrows():
    mid = row[0]
    movie_name = row[1]
    if mid not in sid2movie_map:
        sid2movie_map[mid] = movie_name
    else:
        continue

api = "http://127.0.0.1:6006/v1/image_text2text/api/"

prompt = """
你是一个专业的电影特征提取专家，能从电影海报或宣传图中提取关键信息，并打算将将信息作用于个性化推荐上。
<电影名称> 
{}

<任务> 
请根据电影名称及其海报或宣传图中提取详细的内容，如主要人物、环境背景、色彩使用及其潜在的情感和主题指向，这将有助于深入分析电影的视觉传达策略和观众吸引力。

<约束> 
1、客观分析时需详细描述海报中可见的所有元素，包括人物、物体、颜色、文字等，并尝试解释它们如何与电影的主题和情感相联系。
2、避免基于个人偏好或预先知识对电影进行评价或批判，分析应基于海报本身提供的视觉信息。
3、禁止揭露电影的关键情节或结局，以防剧透。
4、尽可能使用专业的视觉分析术语和理论框架，支持分析的深度和准确性。
5、请用`中文`回答，并且字数不超过200字。
"""


def request_one_image(sid, movie_name, prompt):
    message_input = {
        "image": f"pic/{sid}.jpg",
        "text": prompt.format(movie_name)
    }
    res = requests.post(api, json=message_input)
    return res.text


## 单个执行
# res = request_one_image(sid=2, movie_name="Jumanji (1995)", prompt=prompt)
# print(res)

## 全部执行
modeled_res = {}
for sid, movie_name in tqdm(sid2movie_map.items()):
    res = request_one_image(sid=sid, movie_name=movie_name, prompt=prompt)
    modeled_res[sid] = res
    # print(res)

with open("LLM_modeled_res.pkl", "wb") as f:
    pickle.dump(modeled_res)
